//===------ YieldOnceCheck.cpp - Check usage of yields in accessors  ------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

// This pass statically verifies that yield-once coroutines, such as the
// generalized accessors `read` and `modify`, yield exactly once in every
// invocation, and diagnoses any violation of this property. This pass uses a
// linear-time, data-flow analysis to check that every path in the control-flow
// graph of the coroutine has a yield instruction before a return instruction.

#define DEBUG_TYPE "yield-once-check"

#include "polarphp/ast/AstWalker.h"
#include "polarphp/ast/DiagnosticsPIL.h"
#include "polarphp/ast/Expr.h"
#include "polarphp/ast/Stmt.h"
#include "polarphp/pil/lang/BasicBlockUtils.h"
#include "polarphp/pil/lang/Dominance.h"
#include "polarphp/pil/optimizer/passmgr/Transforms.h"
#include "llvm/ADT/BreadthFirstIterator.h"
#include "llvm/ADT/DenseSet.h"

using namespace polar;

namespace {

class YieldOnceCheck : public PILFunctionTransform {

   template <typename... T, typename... U>
   static InFlightDiagnostic diagnose(AstContext &Context, SourceLoc loc,
                                      Diag<T...> diag, U &&... args) {
      return Context.Diags.diagnose(loc, diag, std::forward<U>(args)...);
   }

   /// Data-flow analysis state that is associated with basic blocks.
   struct BBState {

      /// Indicates whether a basic block is encountered before seeing a yield
      /// (BeforeYield) or after seeing a yield (AfterYield), or in both states
      /// (Conflict). This enum is a semi-lattice where Conflict is the top and
      /// the merge of BeforeYield and AfterYield states is Conflict.
      enum YieldState { BeforeYield, AfterYield, Conflict } yieldState;

   private:
      // The following states are maintained for emitting diagnostics.

      /// For AfterYield and Conflict states, this field records the yield
      /// instruction that was seen while propagating the state.
      PILInstruction *yieldInst;

      /// For Conflict state, these fields record the basic blocks that
      /// propagated the 'AfterYield' and 'BeforeYield' states which resulted
      /// in the Conflict.
      PILBasicBlock *yieldingPred;
      PILBasicBlock *nonYieldingPred;

      BBState(YieldState yState, PILInstruction *yieldI, PILBasicBlock *yieldPred,
              PILBasicBlock *noYieldPred)
         : yieldState(yState), yieldInst(yieldI), yieldingPred(yieldPred),
           nonYieldingPred(noYieldPred) {}

   public:
      static BBState getInitialState() {
         return BBState(BeforeYield, nullptr, nullptr, nullptr);
      }

      static BBState getAfterYieldState(PILInstruction *yieldI) {
         assert(yieldI);
         return BBState(AfterYield, yieldI, nullptr, nullptr);
      }

      static BBState getConflictState(PILInstruction *yieldI,
                                      PILBasicBlock *yieldPred,
                                      PILBasicBlock *noYieldPred) {
         assert(yieldI && yieldPred && noYieldPred);
         return BBState(Conflict, yieldI, yieldPred, noYieldPred);
      }

      PILInstruction *getYieldInstruction() const {
         assert(yieldState == AfterYield || yieldState == Conflict);
         return yieldInst;
      }

      PILBasicBlock *getYieldingPred() {
         assert(yieldState == Conflict);
         return yieldingPred;
      }

      PILBasicBlock *getNonYieldingPred() {
         assert(yieldState == Conflict);
         return nonYieldingPred;
      }
   };

   /// A structure that records an error found during the analysis along with
   /// some context information that will be used by diagnostics.
   struct YieldError {
      /// The kind of error.
      enum Kind { MultipleYield, ReturnBeforeYield, ReturnOnConflict } errorKind;
      /// The termination instruction where the error should be reported.
      PILInstruction *termInst;
      /// The input state when the error is encountered.
      BBState inState;

   private:
      YieldError(Kind kind, PILInstruction *term, BBState state)
         : errorKind(kind), termInst(term), inState(state) {}

   public:
      static YieldError getMultipleYieldError(YieldInst *yield, BBState state) {
         assert(state.yieldState != BBState::BeforeYield);
         return YieldError(MultipleYield, yield, state);
      }

      static YieldError getReturnBeforeYieldError(ReturnInst *returnI,
                                                  BBState state) {
         assert(state.yieldState == BBState::BeforeYield);
         return YieldError(ReturnBeforeYield, returnI, state);
      }

      static YieldError getReturnOnConflict(ReturnInst *returnI, BBState state) {
         assert(state.yieldState == BBState::Conflict);
         return YieldError(ReturnOnConflict, returnI, state);
      }
   };

   /// Transfer function of the data-flow analysis.
   ///
   /// \param bb Basic block that should be processed
   /// \param inState BBState at the start of the basic block
   /// \param error out parameter that will contain information about
   /// an error that is detected.
   /// \return the state at the exit of the basic block if it can be computed
   /// and None otherwise.
   static Optional<BBState>
   transferStateThroughBasicBlock(PILBasicBlock *bb, BBState inState,
                                  Optional<YieldError> &error) {
      error = None;
      auto *term = bb->getTerminator();

      if (auto *returnInst = dyn_cast<ReturnInst>(term)) {
         if (inState.yieldState == BBState::BeforeYield) {
            error = YieldError::getReturnBeforeYieldError(returnInst, inState);
            return None;
         }

         if (inState.yieldState == BBState::Conflict) {
            error = YieldError::getReturnOnConflict(returnInst, inState);
            return None;
         }
         return inState;
      }

      if (auto *yieldInst = dyn_cast<YieldInst>(term)) {
         if (inState.yieldState != BBState::BeforeYield) {
            error = YieldError::getMultipleYieldError(yieldInst, inState);
            return None;
         }

         // If the current state is BeforeYield and if the basic block ends in a
         // yield the new state is AfterYield.
         return inState.getAfterYieldState(term);
      }

      // We cannot have throws within generalized accessors.
      assert(!isa<ThrowInst>(term));

      return inState;
   }

   /// Merge operation of the data-flow analysis.
   ///
   /// \param mergeBlock the basic block that is reached with two states
   /// \param oldState the previous state at the entry of the basic block
   /// \param newState the current state at the entry of the basic block
   /// \param newStatePred the predecessor of 'mergeBlock' that has propagated
   ///        the 'newState'.
   /// \param bbToStateMap a map from the basic blocks visited by the analysis
   ///        to the BBStates in which they were seen. This is used to identify
   ///        blocks that propagate conflicting states when the merge results
   ///        in a conflict.
   /// \return the new state obtained by merging the oldState with the newState
   static BBState merge(PILBasicBlock *mergeBlock, BBState oldState,
                        BBState newState, PILBasicBlock *newStatePred,
                        llvm::DenseMap<PILBasicBlock *, BBState> &bbToStateMap) {
      auto oldYieldState = oldState.yieldState;
      auto newYieldState = newState.yieldState;

      if (oldYieldState == BBState::Conflict) {
         return oldState;
      }

      if (newYieldState == BBState::Conflict) {
         return newState;
      }

      if (oldYieldState == newYieldState) {
         return oldState;
      }

      // Here, one state is AfterYield and the other one is BeforeYield.
      // Merging them will result in Conflict.
      assert((newYieldState == BBState::AfterYield &&
              oldYieldState == BBState::BeforeYield) ||
             (newYieldState == BBState::BeforeYield &&
              oldYieldState == BBState::AfterYield));

      // For diagnostics, find another predecessor of 'mergeBlock' that was
      // previously seen by the analysis. This predecessor would have
      // propagated the 'oldState'.
      PILBasicBlock *oldStatePred = nullptr;
      for (auto predBB : mergeBlock->getPredecessorBlocks()) {
         if (predBB != newStatePred && bbToStateMap.count(predBB)) {
            oldStatePred = predBB;
            break;
         }
      }
      assert(oldStatePred);

      if (oldState.yieldState == BBState::BeforeYield) {
         return BBState::getConflictState(newState.getYieldInstruction(),
            /* yieldPred */ newStatePred,
            /* noYieldPred */ oldStatePred);
      } else {
         return BBState::getConflictState(oldState.getYieldInstruction(),
            /* yieldPred */ oldStatePred,
            /* noYieldPred */ newStatePred);
      }
   }

   /// Perform a data-flow analysis to check whether there is exactly one
   /// yield before a return in every path in the control-flow graph.
   /// Diagnostics are not reported for nodes unreachable from the entry of
   /// the control-flow graph.
   void diagnoseYieldOnceUsage(PILFunction &fun) {
      llvm::DenseMap<PILBasicBlock *, BBState> bbToStateMap;

      SmallVector<PILBasicBlock *, 16> worklist;

      auto *entryBB = fun.getEntryBlock();
      bbToStateMap.try_emplace(entryBB, BBState::getInitialState());
      worklist.push_back(entryBB);

      // ReturnBeforeYield errors, which denote that no paths yield before
      // returning, are not diagnosed until the analysis completes, in order to
      // distinguish them from ReturnOnConflict errors, which happen when some
      // paths yield and some don't.
      Optional<YieldError> returnBeforeYieldError = None;

      // The algorithm uses a worklist to propagate the state through basic
      // blocks until a fix point. Since the state lattice has height one, each
      // basic block will be visited at most twice, and at most once if there are
      // no conflicts (which are errors). The basic blocks are added to the
      // worklist in a breadth-first fashion. The order of visiting basic blocks
      // is not important for correctness, but it could change the errors
      // diagnosed when there are multiple errors. Breadth-first order diagnoses
      // errors along shorter paths to return.
      while (!worklist.empty()) {
         PILBasicBlock *bb = worklist.pop_back_val();
         const BBState &state = bbToStateMap.find(bb)->getSecond();

         Optional<YieldError> errorResult = None;
         auto resultState = transferStateThroughBasicBlock(bb, state, errorResult);

         if (!resultState.hasValue()) {
            auto error = errorResult.getValue();

            // ReturnBeforeYield errors will not be reported until the analysis
            // completes. So record it and continue.
            if (error.errorKind == YieldError::ReturnBeforeYield) {
               if (!returnBeforeYieldError.hasValue()) {
                  returnBeforeYieldError = error;
               }
               continue;
            }

            emitDiagnostics(error, fun, bbToStateMap);
            return;
         }

         auto nextState = resultState.getValue();

         for (auto *succBB : bb->getSuccessorBlocks()) {
            // Optimistically try to set the state of the successor as next state.
            auto insertResult = bbToStateMap.try_emplace(succBB, nextState);

            // If the insertion was successful, it means we are seeing the successor
            // for the first time. Add the successor to the worklist.
            if (insertResult.second) {
               worklist.insert(worklist.begin(), succBB);
               continue;
            }

            // Here the successor already has a state. Merge the current and
            // previous states and propagate it if it is different from the
            // old state.
            const auto &oldState = insertResult.first->second;
            auto mergedState = merge(succBB, oldState, nextState, bb, bbToStateMap);

            if (mergedState.yieldState == oldState.yieldState)
               continue;

            // Even though at this point there has to be an error since there is an
            // inconsistency between states coming along two different paths,
            // continue propagation of this conflict state to determine
            // whether this results in multiple-yields error or return-on-conflict
            // error.
            insertResult.first->second = mergedState;
            worklist.insert(worklist.begin(), succBB);
         }
      }

      if (returnBeforeYieldError.hasValue()) {
         emitDiagnostics(returnBeforeYieldError.getValue(), fun, bbToStateMap);
      }
   }

   void emitDiagnostics(YieldError &error, PILFunction &fun,
                        llvm::DenseMap<PILBasicBlock *, BBState> &visitedBBs) {
      AstContext &astCtx = fun.getModule().getAstContext();

      switch (error.errorKind) {
         case YieldError::ReturnBeforeYield: {
            diagnose(astCtx, error.termInst->getLoc().getSourceLoc(),
                     diag::return_before_yield);
            return;
         }
         case YieldError::MultipleYield: {
            diagnose(astCtx, error.termInst->getLoc().getSourceLoc(),
                     diag::multiple_yields);

            // Add a note that points to the previous yield.
            diagnose(astCtx,
                     error.inState.getYieldInstruction()->getLoc().getSourceLoc(),
                     diag::previous_yield);
            return;
         }
         case YieldError::ReturnOnConflict: {
            // Emit an error on the return statement.
            diagnose(astCtx, error.termInst->getLoc().getSourceLoc(),
                     diag::possible_return_before_yield);

            // Here, the yield state of 'error' is Conflict.
            auto &conflictState = error.inState;

            // Emit a note that pin points the branch construct that resulted in
            // this conflict. Note that a conflict state is created at a merge block
            // when along one incoming edge the analysis sees a BeforeYield state
            // and along another it sees an AfterYield state.
            // Also note that, by the definition of the merge operation,
            // 'error.yieldingPred()' is the immediate predecessor of the merge block
            // that propagates AfterYield state, and 'error.nonYieldingPred()' is
            // the immediate predecessor of the merge block that propagates a
            // BeforeYield state.
            auto yieldPred = conflictState.getYieldingPred();
            auto noYieldPred = conflictState.getNonYieldingPred();

            // Step 1: find a branching PIL instruction where one branch has
            // 'yieldPred' and another branch has 'noYieldPred'. For instance,
            // in the following example, cond_br is the instruction to find.
            //              cond_br bb1, bb2
            //              bb1:
            //                  yield resume yieldPred, unwind err
            //              bb2:
            //                  br noYieldPred
            //              yieldPred:
            //                  br mergePoint
            //              noYieldPred:
            //                  br mergePoint
            //              mergePoint:
            //                  ...
            // Intuitively, the conflicting branch is the instruction where
            // 'yieldPred' and 'noYieldPred' meet when traversing the CFG in the
            // reverse order. More formally, the branching instruction is a
            // "lowest common ancestor" of 'yieldPred' and 'noYieldPred' in the
            // the DAG obtained from the CFG by ignoring back edges of loops.
            //
            // Note that the lowest common ancestor may not be unique in a DAG.
            // But, any such ancestor could be considered as the conflicting branch as
            // 'yieldPred' and 'noYieldPred' will belong to two different branches of
            // every such ancestor.

            // Find all transitive predecessors of 'yieldPred' that were visited
            // during the analysis
            SmallPtrSet<PILBasicBlock *, 8> predecessorsOfYieldPred;
            for (auto *predBB :
               llvm::breadth_first<llvm::Inverse<PILBasicBlock *>>(yieldPred)) {
               if (visitedBBs.count(predBB)) {
                  predecessorsOfYieldPred.insert(predBB);
               }
            }

            // Find the first predecessor of 'noYieldPred' that is also a predecessor
            // of 'yieldPred', in the breadth-first search order of the reversed CFG.
            PILBasicBlock *lowestCommonAncestorBB = nullptr;
            SmallPtrSet<PILBasicBlock *, 8> predecessorsOfNoYieldPred;
            for (auto *pred :
               llvm::breadth_first<llvm::Inverse<PILBasicBlock *>>(noYieldPred)) {
               if (!visitedBBs.count(pred)) {
                  continue;
               }
               if (predecessorsOfYieldPred.count(pred)) {
                  lowestCommonAncestorBB = pred;
                  break;
               }
               predecessorsOfNoYieldPred.insert(pred);
            }
            assert(lowestCommonAncestorBB);

            auto *conflictingBranch = lowestCommonAncestorBB->getTerminator();

            // Step 2: Find the target basic block of the 'conflictingBranch' that
            // doesn't yield.
            PILBasicBlock *noYieldTarget = nullptr;
            for (auto *succ : lowestCommonAncestorBB->getSuccessorBlocks()) {
               if (predecessorsOfNoYieldPred.count(succ)) {
                  noYieldTarget = succ;
                  break;
               }
            }
            assert(noYieldTarget);

            // Step 3: Report specialized diagnostics for each kind of conflicting
            // branch.

            // For conditional-branch instructions, which correspond to the
            // conditions of 'if', 'where' or 'guard' statements, report for what
            // truth value the branch doesn't yield.
            if (auto *condbr = dyn_cast<CondBranchInst>(conflictingBranch)) {
               diagnose(astCtx, condbr->getLoc().getSourceLoc(),
                        diag::branch_doesnt_yield,
                  /*non-yielding branch*/ condbr->getTrueBB() == noYieldTarget);
               return;
            }

            // For switch_enum instructions, report the case that doesn't yield.
            enum SwitchCaseKind { Default, OptionNil, OptionSome };

            if (auto *switchEnum = dyn_cast<SwitchEnumInstBase>(conflictingBranch)) {
               auto enumCaseLoc = noYieldTarget->begin()->getLoc().getSourceLoc();

               if (switchEnum->hasDefault() &&
                   switchEnum->getDefaultBB() == noYieldTarget) {
                  diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, Default);
                  return;
               }

               // Find the case identifier that doesn't yield.
               NullablePtr<EnumElementDecl> enumElemDecl =
                  switchEnum->getUniqueCaseForDestination(noYieldTarget);
               assert(enumElemDecl.isNonNull());

               // Specialize diagnostics for cases of an optional.
               if (enumElemDecl.get() == astCtx.getOptionalNoneDecl()) {
                  diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, OptionNil);
               } else if (enumElemDecl.get() == astCtx.getOptionalSomeDecl()) {
                  diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, OptionSome);
               } else {
                  diagnose(astCtx, enumCaseLoc, diag::named_case_doesnt_yield,
                           enumElemDecl.get()->getName());
               }
               return;
            }

            // For switch_value instructions, report the case number that doesn't
            // yield.
            if (auto *switchValue = dyn_cast<SwitchValueInst>(conflictingBranch)) {
               auto enumCaseLoc = noYieldTarget->begin()->getLoc().getSourceLoc();

               if (switchValue->hasDefault() &&
                   switchValue->getDefaultBB() == noYieldTarget) {
                  diagnose(astCtx, enumCaseLoc, diag::case_doesnt_yield, Default);
                  return;
               }
               // Find the case that doesn't yield.
               Optional<unsigned> caseNumberOpt =
                  switchValue->getUniqueCaseForDestination(noYieldTarget);
               assert(caseNumberOpt.hasValue());

               auto caseNumber = caseNumberOpt.getValue();
               diagnose(
                  astCtx, enumCaseLoc, diag::switch_value_case_doesnt_yield,
                  (Twine(caseNumber) + llvm::getOrdinalSuffix(caseNumber)).str());
               return;
            }

            // For try-apply instructions, report whether throwing or non-throwing
            // case doesn't yield.
            if (auto *tryApply = dyn_cast<TryApplyInst>(conflictingBranch)) {
               diagnose(astCtx, tryApply->getLoc().getSourceLoc(),
                        diag::try_branch_doesnt_yield,
                  /*does error case not yield?*/ tryApply->getErrorBB() ==
                                                 noYieldTarget);
               return;
            }

            llvm_unreachable("unexpected branch resulting in conflicting yield "
                             "states found in generalized accessor");
         }
      }
   }

   /// The entry point to the transformation.
   void run() override {
      auto *fun = getFunction();

      if (fun->getLoweredFunctionType()->getCoroutineKind() !=
          PILCoroutineKind::YieldOnce)
         return;

      diagnoseYieldOnceUsage(*fun);
   }
};

} // end anonymous namespace

PILTransform *polar::createYieldOnceCheck() {
   return new YieldOnceCheck();
}
